package dmp.frame.java.agent.transformer.impl;

import dmp.frame.java.agent.transformer.MyTransformer;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayInputStream;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.reflect.Modifier;
import java.security.ProtectionDomain;
import java.util.HashSet;
import java.util.Set;

public class RibbonRequestTransformer implements MyTransformer {
    private static final Logger log = LoggerFactory.getLogger(RibbonRequestTransformer.class);

    public static boolean isNative(CtMethod method) {
        return Modifier.isNative(method.getModifiers());
    }

    @Override
    public Set<String> getNames() {
        Set<String> set = new HashSet<String>();
        //发起feign增强
        set.add("feign/SynchronousMethodHandler");
        //接收feign增强
        set.add("org/springframework/web/client/RestTemplate");

        return set;
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        log.info("registry Feign request transformer");
        if (className.equals("feign/SynchronousMethodHandler")) {
            return resetFeignRequestTemplate(className, classfileBuffer);
        } else if (className.equals("org/springframework/web/client/RestTemplate")) {
            return resetRestTemplate(className, classfileBuffer);
        }
        return  classfileBuffer;
    }

    public static byte[] resetFeignRequestTemplate(String className,byte[] classfileBuffer){
        String traceName = "1.0.2";
        System.out.println(className);
        try {
            ClassPool cp = ClassPool.getDefault();
            CtClass ctClass = cp.makeClass(new ByteArrayInputStream(classfileBuffer));
            CtMethod ctMethod = ctClass.getDeclaredMethod("executeAndDecode");
            ctMethod.insertBefore("String mdcTraceId = org.slf4j.MDC.get(\""+traceName+"\");\n" +
                    "        if (mdcTraceId==null){\n" +
                    "            mdcTraceId = java.util.UUID.randomUUID().toString().replaceAll(\"-\",\"\");\n" +
                    "        }\n" +
                    "        template.header(\""+traceName+"\", new String[]{mdcTraceId});");
            return ctClass.toBytecode();
        }catch (Throwable e){
            e.printStackTrace();
            System.out.println("loading FeignRequestTemplate fail...");
            return classfileBuffer;
        }
    }

    public static byte[] resetRestTemplate(String className,byte[] classfileBuffer){
        String traceName = "1.0.2";

        System.out.println(className);
        try {
            ClassPool cp = ClassPool.getDefault();
            CtClass ctClass = cp.makeClass(new ByteArrayInputStream(classfileBuffer));
            CtMethod ctMethod = ctClass.getDeclaredMethod("doExecute");
            ctMethod.setBody("{" +
                    "\t\torg.springframework.util.Assert.notNull($1, \"URI is required\");\n" +
                    "\t\torg.springframework.util.Assert.notNull($2, \"HttpMethod is required\");\n" +
                    "\t\torg.springframework.http.client.ClientHttpResponse response = null;\n" +
                    "\t\ttry {\n" +
                    "\t\t\torg.springframework.http.client.ClientHttpRequest request = createRequest($1, $2);\n" +
                    "\t\t\tif ($3 != null) {\n" +
                    "\t\t\t\t$3.doWithRequest(request);\n" +
                    "\t\t\t}\n" +
                    //-----加上"+traceName+"-----
                    "\t\t\torg.springframework.http.HttpHeaders headers = request.getHeaders();\n" +
                    "\t\t\tString mdcTraceId = headers.getFirst(\""+traceName+"\");\n" +
                    "\t\t\tif (mdcTraceId==null){\n" +
                    "\t\t\t\tmdcTraceId = org.slf4j.MDC.get(\""+traceName+"\");\n" +
                    "\t\t\t\tif (mdcTraceId==null){\n" +
                    "\t\t\t\t\tmdcTraceId = java.util.UUID.randomUUID().toString().replaceAll(\"-\",\"\");\n" +
                    "\t\t\t\t}\n" +
                    "\t\t\t}\n" +
                    "\t\t\tString tmp = mdcTraceId;\n" +
                    "\t\t\torg.slf4j.MDC.put(\""+traceName+"\",tmp);\n" +
                    "\t\t\theaders.add(\""+traceName+"\", mdcTraceId);"+
                    //=====加上"+traceName+"=====
                    "\t\t\tresponse = request.execute();\n" +
                    "\t\t\thandleResponse($1, $2, response);\n" +
                    "\t\t\treturn ($4 != null ? $4.extractData(response) : null);\n" +
                    "\t\t}\n" +
                    "\t\tcatch (java.io.IOException ex) {\n" +
                    "\t\t\tString resource = $1.toString();\n" +
                    "\t\t\tString query = $1.getRawQuery();\n" +
                    "\t\t\tresource = (query != null ? resource.substring(0, resource.indexOf('?')) : resource);\n" +
                    "\t\t\tthrow new org.springframework.web.client.ResourceAccessException(\"I/O error on \" + $2.name() +\n" +
                    "\t\t\t\t\t\" request for \\\"\" + resource + \"\\\": \" + ex.getMessage(), ex);\n" +
                    "\t\t}\n" +
                    "\t\tfinally {\n" +
                    "\t\t\tif (response != null) {\n" +
                    "\t\t\t\tresponse.close();\n" +
                    "\t\t\t}\n" +
                    "\t\t}" +
                    "}");
            return ctClass.toBytecode();
        }catch (Throwable e){
            e.printStackTrace();
            System.out.println("loading RestTemplate fail...");
            return classfileBuffer;
        }
    }
}
